import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import cv2
import h5py
from sklearn import preprocessing
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model, load_model, model_from_json
from PIL import Image
import time
import csv
from run_experiment import load_params
import math
import scipy
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split


def get_origin_report(origin_model):
    precision_old = []
    recall_old = []
    f1_old = []
    X_data = []
    Y_data = []
    num_class = 6
    f = csv.reader(open('./data/data_set/BIT_label.csv', 'r', encoding='utf-8'))
    input_shape = (32, 32, 3)  # 3通道图像数据
    for lists in f:
        img = cv2.imread('./data/data_set/BIT/' + lists[1])
        img = cv2.resize(img, (input_shape[0], input_shape[1]))
        X_data.append(img)
        Y_data.append(int(lists[2]))
    train_x, test_x, train_y, test_y = train_test_split(X_data, Y_data, test_size=0.4)
    train_x = np.array(train_x).astype('float32') / 255.
    train_y = np.array(train_y)
    test_x = np.array(test_x).astype('float32') / 255.
    test_y = np.array(test_y)
    train_x_mean = np.mean(train_x, axis=0)
    test_x_mean = np.mean(test_x, axis=0)
    train_x -= train_x_mean
    test_x -= test_x_mean
    lb = preprocessing.LabelBinarizer().fit(np.array(range(num_class)))  # 对标签进行ont_hot编码
    train_y = lb.transform(train_y)  # 因为是多分类任务，必须进行编码处理
    test_y = lb.transform(test_y)
    test_x_perdict = origin_model(test_x)
    test_x_perdict_arg = np.argmax(test_x_perdict, axis=1)
    test_y_arg = np.argmax(test_y, axis=1)
    report_dict = classification_report(test_y_arg, test_x_perdict_arg, output_dict=True)
    report = classification_report(test_y_arg, test_x_perdict_arg)
    print(report)
    for i in np.arange(len(report_dict) - 3):
        precision_old.append(report_dict[str(i)]['precision'])
        recall_old.append(report_dict[str(i)]['recall'])
        f1_old.append(report_dict[str(i)]['f1-score'])
    return precision_old, recall_old, f1_old


def get_end_report(origin_model, model, data):
    precision_new = []
    recall_new = []
    f1_new = []
    grad_path = "./output/grad/" + model + '/' + data + "/generated_samples_low_none_3"
    grad_new_inputs = np.load(os.path.join(grad_path, 'new_inputs.npy'))
    grad_origin_inputs = np.load(os.path.join(grad_path, 'orgin_inputs.npy'))
    grad_new_outputs = np.load(os.path.join(grad_path, 'new_outputs.npy'))
    grad_perdict = origin_model(grad_new_inputs)
    grad_perdict_arg = np.argmax(grad_perdict, axis=1)
    grad_new_outputs_arg = np.argmax(grad_new_outputs, axis=1)
    report_dict = classification_report(grad_new_outputs_arg, grad_perdict_arg, output_dict=True)
    report = classification_report(grad_new_outputs_arg, grad_perdict_arg)
    print(report)
    for i in np.arange(len(report_dict) - 3):
        precision_new.append(report_dict[str(i)]['precision'])
        recall_new.append(report_dict[str(i)]['recall'])
        f1_new.append(report_dict[str(i)]['f1-score'])
    return precision_new, recall_new, f1_new;


def get_rate(precision_old, recall_old, f1_old, precision_new, recall_new, f1_new):
    f1_rate = []
    precision_rate = []
    recall_rate = []
    for old, new in zip(f1_old, f1_new):
        f1_rate.append(old - new)
    for old, new in zip(precision_old, precision_new):
        precision_rate.append(old - new)
    for old, new in zip(recall_old, recall_new):
        recall_rate.append(old - new)
    return precision_rate, recall_rate, f1_rate


def run_evaluate(params):
    params = load_params(params)
    model = params.model
    data = params.data
    model_path = params.model_path
    model_name = data + '_' + model + '.h5'
    origin_model = load_model(os.path.join(model_path, model_name))
    precision_old, recall_old, f1_old = get_origin_report(origin_model)
    print("precision_old: " + str(precision_old))
    print("recall_old: " + str(recall_old))
    print("f1_old: " + str(f1_old))
    print("----------------------------------------------------------------------------------------")
    precision_new, recall_new, f1_new = get_end_report(origin_model, model, data)
    print("precision_new: " + str(precision_new))
    print("recall_new: " + str(recall_new))
    print("f1_new: " + str(f1_new))
    print("----------------------------------------------------------------------------------------")
    precision_rate, recall_rate, f1_rate = get_rate(precision_old, recall_old, f1_old, precision_new, recall_new,
                                                    f1_new)
    print("precision_rate: " + str(precision_rate))
    print("recall_rate: " + str(recall_rate))
    print("f1_rate: " + str(f1_rate))
    print("----------------------------------------------------------------------------------------")
    kill_origin_dict = [
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

    kill_per_dict = [
        [0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1,
         1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0,
         0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1],
        [1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1,
         1, 1, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0,
         0, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1,
         0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0,
         1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1],
        [0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1,
         0, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 1, 1, 0, 0,
         0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0,
         1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1,
         1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0]]

    kill_origin_dict = np.array(kill_origin_dict)
    kill_per_dict = np.array(kill_per_dict)
    kill_origin = []
    for i in np.arange(len(kill_origin_dict)):
        kill_origin.append(kill_origin_dict[i].sum())
    kill_per_class = []
    for i in np.arange(len(kill_per_dict)):
        kill_per_class.append(kill_per_dict[i].sum())
    kill_rate = []
    for old, new in zip(kill_per_class, kill_origin):
        kill_rate.append(old - new)
    print("kill_per_class: " + str(kill_per_class))
    print("kill_origin: " + str(kill_origin))
    print("kill_rate: " + str(kill_rate))
    print("----------------------------------------------------------------------------------------")

    correlation_f1_old_origin, p = scipy.stats.spearmanr(f1_old, kill_origin)
    correlation_precision_old_origin, p = scipy.stats.spearmanr(precision_old, kill_origin)
    correlation_recall_old_origin, p = scipy.stats.spearmanr(recall_old, kill_origin)
    print("correlation_f1_old_origin: " + str(correlation_f1_old_origin))
    print("correlation_precision_old_origin: " + str(correlation_precision_old_origin))
    print("correlation_recall_old_origin: " + str(correlation_recall_old_origin))
    print("----------------------------------------------------------------------------------------")
    correlation_f1_new_origin, p = scipy.stats.spearmanr(f1_new, kill_per_class)
    correlation_precision_new_origin, p = scipy.stats.spearmanr(precision_new, kill_per_class)
    correlation_recall_new_origin, p = scipy.stats.spearmanr(recall_new, kill_per_class)
    print("correlation_f1_new: " + str(correlation_f1_new_origin))
    print("correlation_precision_new: " + str(correlation_precision_new_origin))
    print("correlation_recall_new: " + str(correlation_recall_new_origin))
    print("----------------------------------------------------------------------------------------")
    correlation_f1_rate_origin, p = scipy.stats.spearmanr(f1_rate, kill_rate)
    correlation_precision_rate_origin, p = scipy.stats.spearmanr(precision_rate, kill_rate)
    correlation_recall_rate_origin, p = scipy.stats.spearmanr(recall_rate, kill_rate)
    print("correlation_f1_rate: " + str(correlation_f1_rate_origin))
    print("correlation_precision_rate: " + str(correlation_precision_rate_origin))
    print("correlation_recall_rate: " + str(correlation_recall_rate_origin))
    print("----------------------------------------------------------------------------------------")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Experiments Script For evaluate")
    parser.add_argument("--params_set", nargs='*', type=str, default=["BIT", "DenseNet121"],
                        help="see params folder")
    parser.add_argument("--model", type=str, default="DenseNet121",
                        choices=["INCEPTION_NET", "VGG16", "ResNet50", "DenseNet121"])
    parser.add_argument("--data", type=str, default="BIT", choices=["BKK100", "GTSRB", "BIT", "Car"])
    parser.add_argument("--approach", type=str, default="tarantula")
    parser.add_argument("--model_path", type=str, default="./data/neural_networks")
    parser.add_argument("--deepfault_path", type=str, default="./output/deepfault")
    parser.add_argument("--mutant_path", type=str, default="./output/mutant")
    params = parser.parse_args()
    start_time = time.time()
    run_evaluate(params)
    time_passed_min = (time.time() - start_time) / 60
    print("time passed (minutes): %g" % time_passed_min)
